from dataclasses import dataclass, field
from typing import Dict, List
from .common import Simulator, CostCalculator
from .base_simulator import ColumnParallelLinearSimulator, RowParallelLinearSimulator
from abc import ABC, abstractmethod

class MLPSimulator(Simulator):
    def __init__(self, mp: Dict[str, object]):
        super().__init__()
        self.mp = mp
        self.hidden_size = mp['input_shape'][0][2]
        self.batch_size = mp['input_shape'][0][1]
        self.seq_len = mp['input_shape'][0][0]

    
    def get_calculate_workload(self, tp_size: int) -> List[float]:
        calculator = [0,0,0,0,0]
        per_partition_output_size = 4 * self.hidden_size // tp_size

        num = self.batch_size * self.seq_len * per_partition_output_size
        calculator[2] += 7 * num
        calculator[0] += 2 * num

        return calculator
    
    def get_communication_workload(self, tp_size):
        return 0
    
    def get_dp_communication_workload(self, tp_size, dp_size):
        return 0

@dataclass
class MLPLayer(CostCalculator):
    mlp: MLPSimulator = None
    linear_fc1: ColumnParallelLinearSimulator = None
    linear_fc2: RowParallelLinearSimulator = None


class MLPLayerAbstractBuilder(ABC):
    @abstractmethod
    def build_mlp(self, args: Dict[str, object]): ...
    @abstractmethod
    def build_linear_fc1(self, args: Dict[str, object]): ...
    @abstractmethod
    def build_linear_fc2(self, args: Dict[str, object]): ...
